#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
"""
@File    :   net.py
@Time    :   2024-08-16 20:10
@Author  :   坐公交也用券
@Version :   1.0
@Contact :   faith01238@hotmail.com
@Homepage : https://liumou.site
@Desc    :   虚拟机网卡管理工具
"""
import logging
import sys
from subprocess import getstatusoutput
from sys import exit

# 创建一个自定义的日志记录器
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

# 创建一个自定义的日志处理器，设置其输出格式
formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d | %(funcName)s | %(message)s')
handler = logging.StreamHandler(sys.stdout)
handler.setFormatter(formatter)
logger.addHandler(handler)


class NetDelete:
	def __init__(self):
		self.net_list = []
		self.net_dict = {}
		self.net_type = ""
		self.net_delete_start = 0
		self.net_delete_end = 0
		self.host_id_dict = {}
		self.host_eth_count = 0
		self.id = None
		self.net_count = 0
		self.interface_dict = {}
		self.interface_name = 'vmbr0'

	def get_host_id_list(self):
		"""
		获取主机ID列表
		:return:
		"""
		status, result = getstatusoutput("qm list | awk '{print $1,$2}' | grep -v ^VMID")
		if status == 0:
			for i in result.split("\n"):
				self.host_id_dict[i.split(" ")[0]] = i.split(" ")[1]
		else:
			print("获取主机ID列表失败")
			exit(1)

	def select_host(self):
		"""
		选择需要操作的主机ID
		:return:
		"""
		while True:
			print("请选择需要操作的主机ID：")
			for i in self.host_id_dict:
				print(f"{i} {self.host_id_dict[i]}")
			print("0. 退出")
			choice = input("请输入您的选择：\n")
			if choice == "0":
				exit(0)
			elif choice in self.host_id_dict:
				self.id = choice
				break
			else:
				print("输入错误，请重新输入")

	def get_host_eth_nmber(self):
		"""
		获取主机网卡数量
		:return:
		"""
		cmd = f'qm config {self.id} | grep ^net | wc -l'
		status, net_id = getstatusoutput(cmd)
		if status == 0:
			print(cmd)
			logger.info(f"当前主机{self.id}已存在{net_id}个网卡")
			self.host_eth_count = int(net_id)

	def delete_eth(self):
		"""
		开始添加网卡
		:return:
		"""
		cmd_list = []
		id_ = self.host_eth_count + 1
		logger.info("欲执行的命令如下：")
		for i in range(self.net_delete_start, self.net_delete_end+1):
			cmd = f"qm set {self.id} --delete net{i}"
			print(cmd)
			id_ += 1
			cmd_list.append(cmd)
		# 询问是否开始添加
		while True:
			print("是否开始删除网卡")
			print("1. 开始删除")
			print("2. 退出")
			choice = input("请输入您的选择：\n")
			if choice == "1":
				n_ = 1
				for i in cmd_list:
					status, result = getstatusoutput(i)
					if status == 0:
						logger.info(f"删除第 [{n_} ] 张网卡成功")
					else:
						logger.error(f"删除第 [{n_} ] 张网卡失败")
						print(result)
					n_ += 1
				break
			elif choice == "2":
				exit(0)
				break
			else:
				print("输入错误，请重新输入")

	def set_start(self):
		"""
		设置网卡数量
		:return:
		"""
		while True:
			logger.info("当前主机已存在" + str(self.host_eth_count) + "个网卡")
			try:
				self.net_delete_start = int(input("请设置需要删除的网卡起始ID：\n"))
			except ValueError:
				print("输入错误，请重新输入")
			else:
				break

	def set_end(self):
		"""
		设置网卡数量
		:return:
		"""
		while True:
			logger.info("当前主机已存在" + str(self.host_eth_count) + "个网卡")
			logger.info(f"当前已设置起始网卡ID为：" + str(self.net_delete_start))
			try:
				self.net_delete_end = int(input("请设置需要删除的网卡结束ID：\n"))
			except ValueError:
				print("输入错误，请重新输入")
			else:
				break

	def start(self):
		"""
		开始配置过程，按顺序执行一系列配置步骤。

		此方法无参数，通过调用不同方法实现主机选择、网络类型选择、数量设置、接口列表获取与设置，
		最终完成添加操作。这些步骤是配置过程中的关键环节，确保按照预定顺序执行。
		"""
		self.get_host_id_list()  # 获取主机ID列表
		self.select_host()  # 选择主机
		self.get_host_eth_nmber()
		self.set_start()  # 设置起始
		self.set_end()  # 设置结束
		self.delete_eth()  # 开始删除


class Net:
	def __init__(self):
		self.net_list = []
		self.net_dict = {}
		self.net_type = ""
		self.host_id_dict = {}
		self.host_eth_count = 0
		self.id = None
		self.net_count = 0
		self.interface_dict = {}
		self.interface_name = 'vmbr0'

	def select_net_type(self):
		"""
		选择pve虚拟机网卡类型
		:return:
		"""
		while True:
			print("请选择网卡类型：")
			print("1. [virtio] 默认类型,半虚拟化,拥有最接近物理的性能,但是部分系统无驱动,无法显示接口速率(不影响使用)")
			print("2. [intel e1000] 兼容模式,相比性能会稍微慢,但是兼容性更好,可以显示接口速率(1000M)")
			print("3. [Realtek RTL8139] 兼容模式,兼容性更好,可以显示接口速率(100M)")
			print(
				"4. [intel e1000e] 兼容模式,相比性能会稍微慢,但是兼容性更好,可以显示接口速率(10000M),PVE版本需大于8.1")
			print("0. 退出")
			choice = input("请输入您的选择：\n")
			if choice == "0":
				exit(0)
			elif choice == "1":
				self.net_type = "virtio"
				break
			elif choice == "2":
				self.net_type = "e1000"
				break
			elif choice == "3":
				self.net_type = "rtl8139"
				break
			elif choice == "4":
				self.net_type = "e1000e"
				break
			else:
				print("输入错误，请重新输入")

	def get_host_id_list(self):
		"""
		获取主机ID列表
		:return:
		"""
		status, result = getstatusoutput("qm list | awk '{print $1,$2}' | grep -v ^VMID")
		if status == 0:
			for i in result.split("\n"):
				self.host_id_dict[i.split(" ")[0]] = i.split(" ")[1]
		else:
			print("获取主机ID列表失败")
			exit(1)

	def select_host(self):
		"""
		选择需要操作的主机ID
		:return:
		"""
		while True:
			print("请选择需要操作的主机ID：")
			for i in self.host_id_dict:
				print(f"{i} {self.host_id_dict[i]}")
			print("0. 退出")
			choice = input("请输入您的选择：\n")
			if choice == "0":
				exit(0)
			elif choice in self.host_id_dict:
				self.id = choice
				break
			else:
				print("输入错误，请重新输入")

	def get_host_eth_nmber(self):
		"""
		获取主机网卡数量
		:return:
		"""
		cmd = f'qm config {self.id} | grep ^net | wc -l'
		status, net_id = getstatusoutput(cmd)
		if status == 0:
			print(cmd)
			logger.info(f"当前主机{self.id}已存在{net_id}个网卡")
			self.host_eth_count = int(net_id)

	def add(self):
		"""
		开始添加网卡
		:return:
		"""
		cmd_list = []
		id_ = self.host_eth_count + 1
		logger.info("欲执行的命令如下：")
		for i in range(self.net_count):
			cmd = f"qm set {self.id} --net{id_} {self.net_type},bridge={self.interface_name}"
			print(cmd)
			id_ += 1
			cmd_list.append(cmd)
		# 询问是否开始添加
		while True:
			print("是否开始添加网卡？")
			print("1. 开始添加")
			print("2. 退出")
			choice = input("请输入您的选择：\n")
			if choice == "1":
				n_ = 1
				for i in cmd_list:
					status, result = getstatusoutput(i)
					if status == 0:
						logger.info(f"添加第 [{n_} ] 张网卡成功")
					else:
						logger.error(f"添加第 [{n_} ] 张网卡失败")
					n_ += 1
			elif choice == "2":
				exit(0)
				break
			else:
				print("输入错误，请重新输入")

	def get_interface_list(self):
		"""
		获取网卡接口列表
		:return:
		"""
		id_ = 1
		status, result = getstatusoutput("ip link | awk '{print $2}' | sed 's/://' | grep vmbr")
		if status == 0:
			for i in result.split("\n"):
				self.interface_dict[id_] = i
				logger.info(f"{id_} {i}")
				id_ += 1
		else:
			print("获取网卡接口列表失败")
			exit(1)

	def set_interface(self):
		"""
		设置网卡接口
		:return:
		"""
		while True:
			print("请设置网卡接口：")
			for i in self.interface_dict:
				print(f"{i} {self.interface_dict[i]}")
			try:
				self.interface_name = self.interface_dict[int(input("请输入您的选择："))]
				break
			except ValueError:
				print("输入错误，请重新输入")

	def set_number(self):
		"""
		设置网卡数量
		:return:
		"""
		while True:
			logger.info("当前主机已存在" + str(self.host_eth_count) + "个网卡")
			print("请设置需要新增的网卡数量：")
			try:
				self.net_count = int(input("请设置网卡数量：\n"))
			except ValueError:
				print("输入错误，请重新输入")
			else:
				break

	def start(self):
		"""
		开始配置过程，按顺序执行一系列配置步骤。

		此方法无参数，通过调用不同方法实现主机选择、网络类型选择、数量设置、接口列表获取与设置，
		最终完成添加操作。这些步骤是配置过程中的关键环节，确保按照预定顺序执行。
		"""
		self.get_host_id_list()  # 获取主机ID列表
		self.select_host()  # 选择主机
		self.get_host_eth_nmber()
		self.select_net_type()  # 选择网络类型
		self.set_number()  # 设置数量
		self.get_interface_list()  # 获取接口列表
		self.set_interface()  # 设置接口
		self.add()  # 添加配置


if __name__ == '__main__':
	while True:
		choice = input("请选择要执行的操作：\n1. 添加网卡\n2. 删除网卡\n0. 退出\n")
		if choice == "1":
			logger.info("添加网卡")
			n = Net()
			n.start()
		elif choice == "2":
			logger.info("删除网卡")
			n = NetDelete()
			n.start()
		elif choice == "0":
			exit(0)
		else:
			logger.warning("输入错误，请重新输入")

